import requests
import json
import time
from langchain_core.embeddings import Embeddings
from typing import List,Optional
from langchain.llms.base import LLM
from langchain.llms.utils import enforce_stop_tokens

API_KEY = 'sk-b3e7b4af954081fbbdf5342e6f63ca6f'

URL = 'https://api.baichuan-ai.com/v1/chat/completions'

EMBEDDING_URL = 'http://api.baichuan-ai.com/v1/embeddings'

headers = {
    'Content-Type': 'application/json',
    'Authorization': 'Bearer ' + API_KEY
}

def send_message(prompt, json_flag=False):
    data = {
        'model': 'Baichuan4-Air',
        'messages': prompt,
        'temperature': 0.3,
        'top_p': 0.85,
        'max_tokens': 2048,
        'with_search_enhance': False
    }
    response = requests.post(URL, headers=headers, json=data)
    while response.status_code == 429:
        print('API rate limit exceeded, waiting 10 seconds...')
        time.sleep(5)
        response = requests.post(URL, headers=headers, json=data)
    if response.status_code == 200:
        result = json.loads(response.text)
        if result['choices'][0]['finish_reason'] == 'content_filter':
            print('Content filtered, retrying...')
            time.sleep(5)
            response = requests.post(URL, headers=headers, json=data)
            result = json.loads(response.text)
        tmp_content = result['choices'][0]['message']['content'].strip()
        #time.sleep(5)
        return tmp_content
    else:
        return None

def send_embedding(text_list):
    data = {
        'model': 'Baichuan-Text-Embedding',
        'input': text_list
    }
    response = requests.post(EMBEDDING_URL, headers=headers, json=data)
    #print(response)
    while response.status_code == 429:
        print('API rate limit exceeded, waiting 10 seconds...')
        time.sleep(5)
        response = requests.post(EMBEDDING_URL, headers=headers, json=data)
    if response.status_code == 200:
        result = json.loads(response.text)
        return [item['embedding'] for item in result['data']]
    else:
        return None

class BaichuanEmbedding(Embeddings):
    
    def __init__(self,send_len):
        #super.__init__()
        self.SEND_LEN = send_len
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed search docs."""
        res_list = []
        cur_list = []
        for item in texts:
            cur_list.append(item)
            if len(cur_list) == self.SEND_LEN:
                res_list += send_embedding(cur_list)
                cur_list = []
                time.sleep(2)
        
        if len(cur_list) > 0:
            res_list += send_embedding(cur_list)
    
        return res_list
    
    def embed_query(self, text: str) -> List[float]:
        """Embed query text."""
        return send_embedding([text])[0]

class BaichuanLLM(LLM):
    def __init__(self):
        super().__init__()

    @property
    def _llm_type(self) -> str:
        return "Baichuan"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        
        response = send_message(prompt=[{"role": "user", "content": prompt}])
        #print(response)
        if not response:
            return "错误"
        
        if stop is not None:
            response = enforce_stop_tokens(response, stop)
        #self.history = self.history+[[None, resp['response']]]
        return response

